package com.kingsoft.dc.khaos.module.spark.util;

import org.apache.hadoop.security.UserGroupInformation;
import org.apache.spark.SparkFiles;
import org.elasticsearch.hadoop.EsHadoopException;
import org.elasticsearch.hadoop.cfg.ConfigurationOptions;
import org.elasticsearch.hadoop.mr.security.HadoopUser;
import org.elasticsearch.hadoop.security.User;
import org.elasticsearch.hadoop.security.UserProvider;
import org.elasticsearch.hadoop.cfg.Settings;
import org.apache.hadoop.conf.Configuration;

import java.io.File;
import java.io.IOException;

/**
 * Create by wuxiang on 2022/5/9
 */
public class ESUserProvider extends UserProvider {
    public static final String SPARK_ES_PRINCIPAL = "es.spark.principal";
    public static final String SPARK_ES_KEYTAB = "es.spark.keytab";
    public static final String ES_NET_SPNEGO_AUTH_ELASTICSEARCH_PRINCIPAL = "es.net.spnego.auth.elasticsearch.principal";

    /**
     * @deprecated Use {@link UserProvider#create(Settings)} instead, or if it is important that you
     * explicitly create a HadoopUserProvider, then use {@link ESUserProvider#create(Settings)}
     * instead
     */
    public ESUserProvider() {
        super();
    }

    public static ESUserProvider create(Settings settings) {

        ESUserProvider provider = new ESUserProvider();
        provider.setSettings(settings);
        return provider;
    }

    @Override
    public User getUser() {
        try {
            UserGroupInformation ugi = UserGroupInformation.getLoginUser();
            if (ugi == null || !ugi.hasKerberosCredentials())
                ugi = doLogin(settings);
            return new HadoopUser(ugi, getSettings());
        } catch (IOException e) {
            throw new EsHadoopException("Could not retrieve the current user", e);
        }
    }

    public UserGroupInformation doLogin(Settings settings) throws IOException {
        String keytab = settings.getProperty(SPARK_ES_KEYTAB);
        String keytabPath = SparkFiles.get(keytab);
        String principal = settings.getProperty(SPARK_ES_PRINCIPAL);
        String spnegoPrincipal = settings.getProperty(ES_NET_SPNEGO_AUTH_ELASTICSEARCH_PRINCIPAL);

        Configuration krb5_conf = new Configuration();
        krb5_conf.set("hadoop.security.authentication", "kerberos");
        krb5_conf.set(ConfigurationOptions.ES_NET_SPNEGO_AUTH_ELASTICSEARCH_PRINCIPAL, spnegoPrincipal);
        krb5_conf.set(ConfigurationOptions.ES_SECURITY_AUTHENTICATION, "kerberos");
        //判断文件是否存在
        if (!new File(keytabPath).exists())
            throw new RuntimeException("executor端登录失败， keytab文件(" + keytabPath + ")不存在,请通过sparkContext.addFile将文件加入到executor节点");
        if (!new File("/etc/krb5.conf").exists()) throw new RuntimeException("executor端登录失败，/etc/krb5.conf 文件不存在");
        return Krb5Utils.doLogin(principal, keytabPath, krb5_conf);
    }

}
